import codecs
import random


random.seed(222)


prompt_path = '/home/rickwwang/project_own/story_generation/data/test.wp_target_500.format.prompt'
story_path = '/home/rickwwang/project_own/story_generation/data/test.wp_target_500.format.story'

# for seq2seq format, only story
result_src = '/home/rickwwang/project_own/story_generation/data/evaluate/prompt_rank_src.txt'
result_tgt = '/home/rickwwang/project_own/story_generation/data/evaluate/prompt_rank_tgt.txt'

# for lm format, prompt + story
result_lm = '/home/rickwwang/project_own/story_generation/data/evaluate/prompt_rank_src_tgt.txt'

all_prompt = codecs.open(prompt_path, 'r', encoding='utf8').readlines()
all_story = codecs.open(story_path, 'r', encoding='utf8').readlines()

prompt_set = list(set(map(lambda x: x.strip(), all_prompt)))

# select 1000 stories
all_zip = zip(all_prompt, all_story)
random.shuffle(all_zip)
all_zip = all_zip[:1000]

srcs = []
tgts = []
lms = []

for p, s in all_zip:
    p = p.strip()
    s = s.strip()
    srcs.append(p)
    tgts.append(s)
    lms.append(' <SEP> '.join([p, s]))

    shuffle_ps = random.sample(prompt_set, 9)
    for shuffle_p in shuffle_ps:
        while shuffle_p == p:
            shuffle_p = random.choice(prompt_set)

        srcs.append(shuffle_p)
        tgts.append(s)
        lms.append(' <SEP> '.join([shuffle_p, s]))


with codecs.open(result_src, 'w', encoding='utf8') as fout:
    fout.write('\n'.join(srcs))
    fout.write('\n')
with codecs.open(result_tgt, 'w', encoding='utf8') as fout:
    fout.write('\n'.join(tgts))
    fout.write('\n')
with codecs.open(result_lm, 'w', encoding='utf8') as fout:
    fout.write('\n'.join(lms))
    fout.write('\n')
